package com.k8.cluster.clusterInvokerImpl;


import com.k8.cluster.loadBalanceImpl.AbstractConsistentHashLoadBalance;
import com.k8.rpc.cluster.ClusterInvoker;
import com.k8.rpc.cluster.loadBalance.LoadBalance;
import com.k8.rpc.invoke.Invoker;
import com.k8.rpc.invoke.invocation.Invocation;
import com.k8.rpc.result.Result;

/**
 * @Author: k8
 * @CreateTime: 2025-03-20
 * @Version: 1.0
 */
public abstract class AbstractSessionSupportClusterInvoker<T> implements ClusterInvoker<T> {
    public static InheritableThreadLocal<String> sessionKey = new InheritableThreadLocal<>();
    public static InheritableThreadLocal<Boolean> openSession = new InheritableThreadLocal<>();

    protected LoadBalance<T> loadBalance;

    @Override
    public void setLoadBalance(LoadBalance<T> loadBalance) {
        this.loadBalance = loadBalance;
    }

    @Override
    public void destroy() {
    }

    @Override
    public Result invoke(Invocation invocation) throws Throwable {
        if (openSession.get() == null || !openSession.get()) {
            return doInvoke(invocation);
        }
        if (loadBalance instanceof AbstractConsistentHashLoadBalance<T> abstractConsistentHashLoadBalance) {
            String key = sessionKey.get();
            if (key == null) {
                return doInvoke(invocation);
            }
            Invoker<T> invoker = abstractConsistentHashLoadBalance.getInvoker(key);
            return invoker.invoke(invocation);
        } else {
            return doInvoke(invocation);
        }
    }

    abstract Result doInvoke(Invocation invocation) throws Throwable;
}
